package fem2.strategies;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.SolveFailedException;

import java.util.ArrayList;

import math2.MathUtilities;
import fem2.Debugger;
import fem2.Echoer;
import fem2.Model;
import fem2.Observer;
import fem2.enu.ConvergenceCriteria;
import fem2.enu.EchoLevelType;
import fem2.enu.PETScKSPType;
import fem2.enu.PETScPrecondType;
import fem2.enu.SolverType;
import fem2.jlinalg.BaitschSolver;
import fem2.jlinalg.MKLLSESolver;
import fem2.jlinalg.PETScLSESolver;
import fem2.jlinalg.Solver;
import fem2.jlinalg.SymbolicTable;
import fem2.jlinalg.UmfPackLSESolver;

/**
 * abstract class for all strategies
 * 
 * @author hbui
 * 
 */
public abstract class Strategy extends Echoer implements Observer {

	protected int MAX_NT = 50;

	protected int nt = 1;
	protected Model model;
	private SolverType solverType = SolverType.MKL_SOLVER;
	protected ArrayList<Observer> observers = new ArrayList<Observer>();
	protected double tol = MathUtilities.tol6;
	protected ConvergenceCriteria convergenceCriteria = ConvergenceCriteria.RELATIVE_RESIDUAL;

	public Strategy(Model m) {
		model = m;
	}

	/**
	 * set the solver type
	 * 
	 * @param type
	 */
	public void setSolverType(SolverType type) {
		solverType = type;
	}

	/**
	 * 
	 * @return the type of solver
	 */
	public SolverType getSolverType() {
		return solverType;
	}

	/**
	 * set the convergence criteria for the strategy
	 * 
	 * @param convergenceCriteria
	 */
	public void setConvergenceCriteria(ConvergenceCriteria convergenceCriteria) {
		this.convergenceCriteria = convergenceCriteria;
	}

	/**
	 * add observer
	 * 
	 * @param o
	 * @return this observer
	 */
	public Observer addObserver(Observer o) {
		Debugger.watch("Observer " + o.getClass().getName() + " added sucessfully");
		observers.add(o);
		return o;
	}

	/**
	 * create the solver based on specified solver type
	 * 
	 * @return
	 */
	public Solver createSolver() {
		switch (getSolverType()) {
		case BAITSCH_SOLVER_GENERIC:
			return new BaitschSolver(SolverType.BAITSCH_SOLVER_GENERIC);
		case BAITSCH_SOLVER_SPOOLES:
			return new BaitschSolver(SolverType.BAITSCH_SOLVER_SPOOLES);
		case UMFPACK_SOLVER:
			return new UmfPackLSESolver();
		case MKL_SOLVER:
			return new MKLLSESolver();
		case PETSC_SOLVER:
			return new PETScLSESolver();
		default:
			throw new Error("invalid solver type");
		}
	}

	/**
	 * initialize the solver. By default, the solver is initialized with size of
	 * model, However, for some method, the size may e larger than size of
	 * model. Therefore, override this method is necessary
	 * 
	 * @param solver
	 * @param n
	 */
	protected void initializeSolver(Solver solver) {
		int n = model.getSize();
		if (solver instanceof BaitschSolver) {
			solver.setSize(n);
		} else if (solver instanceof UmfPackLSESolver || solver instanceof MKLLSESolver) {
			SymbolicTable t = new SymbolicTable(n, n);
			t.initializeFrom(model);
			solver.setSize(t);
		} else if (solver instanceof PETScLSESolver) {
			PETScLSESolver petScLSESolver = (PETScLSESolver) solver;
			SymbolicTable t = new SymbolicTable(n, n);
			t.initializeFrom(model);
			int nz = t.getNumbeOfEntries();
			int[] nnz = t.getNumbeOfRowEntriesAsArray();
			petScLSESolver.setSize(n, nz, nnz);
			petScLSESolver.setKSP(PETScKSPType.KSPGMRES);
			petScLSESolver.setPC(PETScPrecondType.PCILU);
		} else {
			throw new Error("invalid solver");
		}
	}

	/**
	 * set number of time step/step for analysis
	 * 
	 * @param nt
	 */
	public void setNumberOfTimeStep(int nt) {
		this.nt = nt;
	}

	/**
	 * set the tolerance for convergence criteria
	 * 
	 * @param tol
	 */
	public void setTolerance(double tol) {
		this.tol = tol;
	}

	/**
	 * set maximum iteration allowed before convergence
	 * 
	 * @param nt
	 */
	public void setMaxNt(int nt) {
		this.MAX_NT = nt;
	}

	/**
	 * get the current solution vector of the strategy
	 * 
	 * @return
	 */
	public abstract double[] getCurrentUHat();

	/**
	 * get the current first derivative of solution vector
	 * 
	 * @return
	 */
	public abstract double[] getCurrentUdHat();

	/**
	 * get the current second derivative of solution vector
	 * 
	 * @return
	 */
	public abstract double[] getCurrentUddHat();

	/**
	 * get current time step of the strategy
	 * 
	 * @return
	 */
	public abstract double getCurrentTimeStep();

	/**
	 * to be called before solving
	 * 
	 * @param leftHandSide
	 * @param rightHandSide
	 * @param u
	 */
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		model.Initialize();
	}

	/**
	 * to be called by solve at initial of each time step
	 * 
	 * @param leftHandSide
	 * @param rightHandSide
	 * @param u
	 */
	public abstract void timeStepStarted(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u);

	/**
	 * to be called at initial of internal Newton Raphson loop
	 * 
	 * @param leftHandSide
	 * @param rightHandSide
	 * @param u
	 */
	public abstract void iterationStarted(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u);

	/**
	 * to be called at later of internal Newton Raphson loop
	 * 
	 * @param leftHandSide
	 * @param rightHandSide
	 * @param u
	 * @return convergence criteria
	 */
	public abstract boolean iterationFinished(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u);

	/**
	 * to be called by solve at later of each time step
	 * 
	 * @param leftHandSide
	 * @param rightHandSide
	 * @param u
	 */
	public abstract void timeStepFinished(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u);

	/**
	 * to be called after solving complete
	 * 
	 * @param leftHandSideSolver
	 * @param rightHandSide
	 * @param u
	 */
	public void finished(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * finalize the solver
		 */
		leftHandSideSolver.finalize();
	}

	/**
	 * solve the analysis problem
	 * 
	 * @throws SolveFailedException
	 */
	public void solve() throws SolveFailedException {

		Solver solver = createSolver();

		Debugger.warn(solver.getClass().getName());

		solver.setEchoLevel(getEchoLevel());
		ArrayVector rhs = new ArrayVector();
		ArrayVector u = new ArrayVector();

		Debugger.watch("Time stepping start");

		started(solver, rhs, u);
		fireStarted();

		for (int i = 0; i < nt; i++) {

			boolean conv = solveOneStep(i, solver, rhs, u);

			if (!conv) {
				Debugger.warn("Analysis stopped at time step " + (i + 1));
				break;
			}
		}

		finished(solver, rhs, u);
		fireFinished();

		Debugger.watch("Time stepping completed");
	}

	protected boolean solveOneStep(int i, Solver solver, ArrayVector rhs, ArrayVector u) {
		timeStepStarted(solver, rhs, u);
		fireTimeStepStarted();

		boolean conv_criteria = false;

		int cnt = 1;
		while (!conv_criteria) {

			iterationStarted(solver, rhs, u);
			fireIterationStarted();

			conv_criteria = iterationFinished(solver, rhs, u);
			fireIterationFinished();

			if (getEchoLevel().contain(EchoLevelType.OUTPUT)) {
				Debugger.watch("---------substep " + cnt + " completed---------");
			}

			cnt++;
			if (cnt > MAX_NT) {
				Debugger.warn("timestep " + (i + 1) + " does not converge in " + MAX_NT + " loops");
				return false;
			}
		}

		timeStepFinished(solver, rhs, u);
		fireTimeStepFinished();

		if (getEchoLevel().contain(EchoLevelType.OUTPUT)) {
			Debugger.watch("-------------------------timestep " + (i + 1)
					+ " completed-------------------------");
			// Debugger.watch("----------------------------------------------------------------------------------------------");
		}
		return true;
	}

	@Override
	public void fireStarted() {
		for (int i = 0; i < observers.size(); i++) {
			observers.get(i).fireStarted();
		}
	}

	@Override
	public void fireTimeStepStarted() {
		for (int i = 0; i < observers.size(); i++) {
			observers.get(i).fireTimeStepStarted();
		}
	}

	@Override
	public void fireIterationStarted() {
		for (int i = 0; i < observers.size(); i++) {
			observers.get(i).fireIterationStarted();
		}
	}

	@Override
	public void fireIterationFinished() {
		for (int i = 0; i < observers.size(); i++) {
			observers.get(i).fireIterationFinished();
		}
	}

	@Override
	public void fireTimeStepFinished() {
		for (int i = 0; i < observers.size(); i++) {
			observers.get(i).fireTimeStepFinished();
		}
	}

	@Override
	public void fireFinished() {
		for (int i = 0; i < observers.size(); i++) {
			observers.get(i).fireFinished();
		}
	}
}
